import os
from typing import Any, Dict

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI

from .BaseLLM import BaseLLM

# 加载环境变量
load_dotenv()

# 支持的模型列表
MODEL_LIST = ["qwen-turbo", "qwen-plus", "qwen-max", "qwq-plus", ""]


class DashScopeLLM(BaseLLM):
    """
    DashScope LLM实现
    使用OpenAI兼容的API接口连接到阿里云百炼平台
    """

    def __init__(self, model_name: str = "qwen-turbo", **kwargs):
        # 验证模型名称
        if model_name not in MODEL_LIST:
            raise ValueError(f"不支持的模型: {model_name}. 可用模型: {', '.join(MODEL_LIST)}")

        super().__init__(model_name=model_name, **kwargs)
        self.api_key = kwargs.get('api_key', os.getenv("DASHSCOPE_API_KEY"))
        self.base_url = kwargs.get('base_url', 'https://dashscope.aliyuncs.com/compatible-mode/v1')
        self.temperature = kwargs.get('temperature', 0.2)
        self.max_tokens = kwargs.get('max_tokens', 4096)
        self.top_p = kwargs.get('top_p', 0.2)
        self.enable_thinking = kwargs.get('enable_thinking', False)

    def create_llm(self, streaming: bool = False, callbacks=None, **kwargs):
        """
        创建DashScope LLM实例
        
        Args:
            streaming: 是否启用流式输出
            callbacks: 回调函数列表
            **kwargs: 其他配置参数
            
        Returns:
            ChatOpenAI: DashScope LLM实例
        """
        # qwq-plus 模型只支持流式模式
        if self.model_name == 'qwq-plus' and not streaming:
            streaming = True

        config = {
            'model': self.model_name,
            'api_key': self.api_key,
            'base_url': self.base_url,
            'max_retries': kwargs.get('max_retries', 2),
            'streaming': streaming,
            'temperature': kwargs.get('temperature', self.temperature),
            'max_tokens': kwargs.get('max_tokens', self.max_tokens),
            'top_p': kwargs.get('top_p', self.top_p),  # 直接设置而不是放在 model_kwargs 中
        }

        # 设置 extra_body 参数
        extra_body = kwargs.get('extra_body', {})
        if self.enable_thinking or kwargs.get('enable_thinking', False):
            extra_body['enable_thinking'] = True
        else:
            extra_body['enable_thinking'] = False

        if extra_body:
            config['extra_body'] = extra_body

        if callbacks:
            config['callbacks'] = callbacks

        # 合并额外配置
        config.update({k: v for k, v in kwargs.items()
                       if
                       k not in ['temperature', 'max_tokens', 'top_p', 'max_retries', 'extra_body', 'enable_thinking']})

        llm = ChatOpenAI(**config)
        return llm

    def get_model_info(self) -> Dict[str, Any]:
        """
        获取DashScope模型信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        model_descriptions = {
            'qwen-turbo': '通义千问系列速度最快、成本极低的模型，适合简单任务。',
            'qwen-plus': '能力均衡，推理效果、成本和速度介于通义千问-Max和通义千问-Turbo之间，适合中等复杂任务。',
            'qwen-max': '通义千问系列效果最好的模型，适合复杂、多步骤的任务。',
            'qwq-plus': '基于 Qwen2.5 模型训练的 QwQ 推理模型，通过强化学习大幅度提升了模型推理能力。'
        }

        return {
            'model_name': self.model_name,
            'provider': 'DashScope(阿里云百炼)',
            'type': 'chat',
            'api_base': self.base_url,
            'temperature': self.temperature,
            'max_tokens': self.max_tokens,
            'top_p': self.top_p,
            'enable_thinking': self.enable_thinking,
            'supports_streaming': True,
            'supports_async': True,
            'supports_tools': True,
            'description': model_descriptions.get(self.model_name, f'{self.model_name}是DashScope平台提供的大语言模型')
        }

    @classmethod
    def get_supported_models(cls) -> list:
        """
        获取支持的模型列表
        
        Returns:
            list: 支持的模型列表
        """
        return MODEL_LIST.copy()

    def invoke(self, messages, **kwargs):
        """
        调用LLM进行对话，为 qwq-plus 模型提供特殊处理
        
        Args:
            messages: 输入消息
            **kwargs: 其他参数
            
        Returns:
            str: LLM响应
        """
        # 对于 qwq-plus 模型，如果不是流式调用，则使用流式方式收集完整响应
        if self.model_name == 'qwq-plus':
            response_parts = []
            for chunk in self.stream(messages, **kwargs):
                response_parts.append(chunk)
            return ''.join(response_parts)
        else:
            # 使用父类的默认实现
            return super().invoke(messages, **kwargs)
